import torch

from problems.DTLZ import DTLZ


class DTLZ3(DTLZ):
    def __init__(self, var_dim: int, obj_dim: int, max_fun_eval, kwargs: dict):
        super().__init__(var_dim, obj_dim, max_fun_eval, kwargs)

    def eval_value(self, dec):
        g = super().get_g(dec)
        obj = g * torch.fliplr(torch.cumprod(
            torch.cat((torch.ones(dec.size(0), 1), torch.cos(dec[:, :self.obj_dim - 1] * torch.pi / 2)), dim=1),
            dim=1)) * torch.cat(
            (torch.ones(dec.size(0), 1), torch.sin(torch.fliplr(dec[:, :self.obj_dim - 1]) * torch.pi / 2)), dim=1)
        return obj

    def get_optimal_solutions(self, size):
        return super().get_optimal_solutions1(size)
